# _*_ coding:utf8 _*_

import cv2
import numpy as np
import uuid


# 从图像中切分出字符
def segment_char(use_image, save_segment_image=False):
    """
    输入RGB图像，输出字符所在的坐标列表，单个元素的格式是：[x1, x2, y1, y2]
    :param use_image:
    :param save_segment_image: 决定是否保存切分后的图片
    :return:
    """
    # 重新调整图像大小，保证是标准
    # 插值方法，像素区域关系进行重采样
    resize_image = cv2.resize(use_image, (150, 50), interpolation=cv2.INTER_AREA)

    image_height = resize_image.shape[0]  # 行数
    image_width = resize_image.shape[1]  # 列数

    # 灰度化和二值化
    source_gray = cv2.cvtColor(resize_image, cv2.COLOR_BGR2GRAY)
    ret_one, threshold_one = cv2.threshold(source_gray, 127, 255, cv2.THRESH_BINARY)

    contours, hierarchy = cv2.findContours(threshold_one, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)

    # 画出轮廓点
    cv2.drawContours(resize_image, contours, -1, (0, 0, 255), 1)
    cv2.imshow('img', resize_image)
    cv2.waitKey(0)

    # 过滤最大的边框
    border_x = 10
    border_y = 5
    min_x = border_x
    min_y = border_y
    max_x = image_width - border_x
    max_y = image_height - border_y

    result_contours = []

    for temp_contour in contours:
        # 开始过滤，超过边界外的都过滤掉
        is_over_border = False
        for i in range(temp_contour.shape[0]):
            if min_x < temp_contour[i, 0, 0] < max_x and min_y < temp_contour[i, 0, 1] < max_y:
                is_over_border = False
            else:
                is_over_border = True
                break
        if not is_over_border:
            result_contours.append(temp_contour)

    # 根据轮廓点计算横向和纵向的直方图
    hist_height = np.zeros(image_height)
    hist_width = np.zeros(image_width)

    # 根据轮廓点统计
    for temp_contour in result_contours:
        for i in range(temp_contour.shape[0]):
            hist_width[temp_contour[i, 0, 0]] += 1
            hist_height[temp_contour[i, 0, 1]] += 1

    # print(hist_height)
    # print(hist_width)

    # 切分数字字符
    # 行取最大值和最小值
    segment_height_list = np.nonzero(hist_height)
    # 上界和下界扩充一个像素
    segment_height = [np.min(segment_height_list) - 1, np.max(segment_height_list) + 1]

    # 列进行切分，遇到为0的位作为分割点
    # 像素阈值
    pixel_threshold = 1

    segment_width = []
    for i in range(1, len(hist_width)):
        # 左界和右界扩充一个像素
        if hist_width[i - 1] <= pixel_threshold < hist_width[i]:
            segment_width.append(i - 1)
        if hist_width[i - 1] > pixel_threshold >= hist_width[i]:
            segment_width.append(i)

    # print(segment_height)
    # print(segment_width)

    # 组合出边框，选取模板
    border_list = []
    expand_x_pixel = 12
    expand_y_pixel = 16

    for i in range(0, len(segment_width), 2):
        # 固定数据
        y_center = (segment_height[0] + segment_height[1]) // 2
        # 可能存在粘连，所以判断下
        # 默认一个字符是25像素
        char_pixel_length = 2 * expand_x_pixel + 1

        # 计算中心点，之后横向扩充
        x_center = (segment_width[i] + segment_width[i + 1]) // 2

        # 需要分割的中心点
        center_list = []

        if 1.5 * char_pixel_length <= (segment_width[i + 1] - segment_width[i]) < 2.5 * char_pixel_length:
            # 分割成两个字符，当前发现最多分割成两个
            x_one_quarter = (segment_width[i] + x_center) // 2
            x_three_quarter = (x_center + segment_width[i + 1]) // 2
            center_list.append(x_one_quarter)
            center_list.append(x_three_quarter)
        else:
            center_list.append(x_center)

        for temp_x_center in center_list:
            # x1, x2, y1, y2
            rect_list = [temp_x_center - expand_x_pixel,
                         temp_x_center + expand_x_pixel,
                         y_center - expand_y_pixel,
                         y_center + expand_y_pixel]
            border_list.append(rect_list)
            cv2.rectangle(resize_image, (rect_list[0], rect_list[2]), (rect_list[1], rect_list[3]), (0, 0, 255), 1)
            print(rect_list)

            if save_segment_image:
                # 保存切分好的图片
                temp_image = threshold_one[rect_list[2]:rect_list[3] + 1, rect_list[0]:rect_list[1] + 1]
                cv2.imwrite('./model_raw_image/' + str(uuid.uuid4()) + '.png', temp_image)

    cv2.imshow('resize_image', resize_image)
    cv2.waitKey(0)

    return border_list


def match_model_digit(char_image):
    """
    根据截取的小块图片，匹配模板，匹配到目标字符
    :param char_image:
    :return:
    """
    char_list = {
        "0": (0, 25),
        "1": (25, 50),
        "2": (50, 75),
        "3": (75, 100),
        "4": (100, 125),
        "5": (125, 150),
        "6": (150, 175),
        "7": (175, 200),
        "8": (200, 225),
        "9": (225, 250)
    }
    # 加载模板图片
    model_image = cv2.imread('model_image/combine_digit.png')
    model_gray = cv2.cvtColor(model_image, cv2.COLOR_BGR2GRAY)

    # 处理比较的图片
    char_gray = cv2.cvtColor(char_image, cv2.COLOR_BGR2GRAY)
    _, char_threshold = cv2.threshold(char_gray, 127, 255, cv2.THRESH_BINARY)

    # 跟模板比较，使用归一化的方差比较，值越小越接近
    maybe_list = cv2.matchTemplate(model_gray, char_threshold, cv2.TM_SQDIFF_NORMED)
    _, _, min_loc, _ = cv2.minMaxLoc(maybe_list)

    # 返回字符
    result_char = ""
    for temp_char, temp_min_max in char_list.items():
        #  - 3 是为了避免在边界的一些值出问题
        if temp_min_max[0] - 3 <= min_loc[0] < temp_min_max[1] - 3:
            result_char = temp_char
            break

    return result_char


def match_model_expression(char_image):
    """
    根据截取的小块图片，匹配模板，匹配到目标字符
    :param char_image:
    :return:
    """
    char_list = {
        "0": (0, 25),
        "1": (25, 50),
        "2": (50, 75),
        "3": (75, 100),
        "4": (100, 125),
        "5": (125, 150),
        "6": (150, 175),
        "7": (175, 200),
        "8": (200, 225),
        "9": (225, 250),
        "+": (250, 275),
        "-": (275, 300),
        "X": (300, 325),
        "=": (325, 350)
    }
    # 加载模板图片
    model_image = cv2.imread('model_image/combine_expression.png')
    model_gray = cv2.cvtColor(model_image, cv2.COLOR_BGR2GRAY)

    # 处理比较的图片
    char_gray = cv2.cvtColor(char_image, cv2.COLOR_BGR2GRAY)
    _, char_threshold = cv2.threshold(char_gray, 127, 255, cv2.THRESH_BINARY)

    # 跟模板比较，使用归一化的方差比较，值越小越接近
    maybe_list = cv2.matchTemplate(model_gray, char_threshold, cv2.TM_SQDIFF_NORMED)
    _, _, min_loc, _ = cv2.minMaxLoc(maybe_list)

    print(min_loc)

    # 返回字符
    result_char = ""
    for temp_char, temp_min_max in char_list.items():
        #  - 3 是为了避免在边界的一些值出问题
        if temp_min_max[0] - 3 <= min_loc[0] < temp_min_max[1] - 3:
            result_char = temp_char
            break

    return result_char


def get_char_list(use_image):
    resize_image = cv2.resize(use_image, (150, 50), interpolation=cv2.INTER_AREA)

    char_location_list = segment_char(resize_image)
    match_char_list = []
    if len(char_location_list) >= 5:
        # 纯数字
        for temp_char_location in char_location_list:
            print(temp_char_location)
            temp_image = resize_image[temp_char_location[2]:temp_char_location[3] + 1,
                         temp_char_location[0]:temp_char_location[1] + 1]
            temp_char = match_model_digit(temp_image)
            match_char_list.append(temp_char)

    else:
        # 数学算式
        for temp_char_location in char_location_list:
            print(temp_char_location)
            temp_image = resize_image[temp_char_location[2]:temp_char_location[3] + 1,
                         temp_char_location[0]:temp_char_location[1] + 1]
            temp_char = match_model_expression(temp_image)
            match_char_list.append(temp_char)

    print(match_char_list)
    return match_char_list


def calculate_digit(char_list):
    """
    计算数据算式，返回指定结果
    :param char_list:
    :return:
    """
    if 3 <= len(char_list) <= 4:
        if char_list[1] == "+":
            return str(int(char_list[0]) + int(char_list[2]))
        elif char_list[1] == "-":
            return str(int(char_list[0]) - int(char_list[2]))
        elif char_list[1] == "X":
            return str(int(char_list[0]) * int(char_list[2]))
    return "".join(char_list)


if __name__ == '__main__':
    image_file = 'test_data/5+9=.png'
    source_image = cv2.imread(image_file)
    recognition_char_list = get_char_list(source_image)
    print(image_file, calculate_digit(recognition_char_list))
